

import cv2
import numpy as np

import glob
import json
#coding : utf-8
import os
import abc
import xml.dom.minidom as xml


class_name=["truck_bed_cover","dirt","empty_bed","back_plate","side_plate"]
categoriesmask = [
    {
        "supercategory": "none",
        "id": 0,
        "name": "truck_bed_cover"
    },
    {
        "supercategory": "none",
        "id": 1,
        "name": "dirt"
    },
    {
        "supercategory": "none",
        "id": 2,
        "name": "empty_bed"
    },
    {
        "supercategory": "none",
        "id": 3,
        "name": "back_plate"
    }
    ,
    {
        "supercategory": "none",
        "id": 4,
        "name": "side_plate"
    }
]
class XmlReader(object):
    __metaclass__ = abc.ABCMeta

    def __init__(self):
        pass

    def read_content(self, filename):
        content = None
        if (False == os.path.exists(filename)):
            return content
        filehandle = None
        try:
            filehandle = open(filename, 'rb')
        except FileNotFoundError as e:
            print(e.strerror)
        try:
            content = filehandle.read()
        except IOError as e:
            print(e.strerror)
        if (None != filehandle):
            filehandle.close()
        if(None != content):
            return content.decode("utf-8", "ignore")
        return content

    @abc.abstractmethod
    def load(self, filename):
        pass


ROOT_PATH ="/home/data"  #os.getcwd()


def put_mask(image, point):

    zeros = np.zeros((image.shape), dtype=np.uint8)
    # zeros_mask = cv2.rectangle(zeros, (bbox[0], bbox[1]), (bbox[2], bbox[3]),
    #                color=(0,0,255), thickness=-1 ) #thickness=-1 表示矩形框内颜色填充
    cv2.fillPoly(zeros, [point], (0, 0, 255), 1)
    gamma = 0

    mask_img = cv2.addWeighted(image, 1, zeros, 0.9, gamma)
    return mask_img


class XmlTester(XmlReader):
    def __init__(self):
        XmlReader.__init__(self)

    def load(self, filename):
        filecontent = XmlReader.read_content(self, filename)
        if None != filecontent:
            # print(filename)
            dom = xml.parseString(filecontent)
            im_size = dom.getElementsByTagName('size')[0]

            im_w = int((im_size.getElementsByTagName(
                'width')[0]).childNodes[0].data)
            im_h = int((im_size.getElementsByTagName(
                "height")[0]).childNodes[0].data)
            self.im_shape = np.array([im_w, im_h])
            self.bbox = []
            self.name = dom.getElementsByTagName(
                'filename')[0].childNodes[0].data
            self.obj_name = []
            for obj in dom.getElementsByTagName('object'):
                box = obj.getElementsByTagName('bndbox')[0]

                b_xmin = int((box.getElementsByTagName(
                    "xmin")[0]).childNodes[0].data)
                b_ymin = int((box.getElementsByTagName(
                    "ymin")[0]).childNodes[0].data)
                b_xmax = int((box.getElementsByTagName(
                    "xmax")[0]).childNodes[0].data)
                b_ymax = int((box.getElementsByTagName(
                    "ymax")[0]).childNodes[0].data)
                b_name = obj.getElementsByTagName(
                    "name")[0].childNodes[0].data
                self.obj_name.append(b_name)
                # print(self.name)
                self.bbox.append([[b_xmin, b_ymin], [b_xmax, b_ymax]])
            self.bbox = np.array(self.bbox)
            self.polygon = []
            for obj in dom.getElementsByTagName('polygon'):
                o_cls = obj.getElementsByTagName('class')[0].childNodes[0].data
                box = obj.getElementsByTagName(
                    'points')[0].childNodes[0].data.replace(";", ",").split(",")
                # for pt in :
                # box = pt.childNodes[0].data.replace(";",",").split(",")
                self.polygon.append(
                    [o_cls, np.array(box).astype(np.int32).reshape(-1,2)])
            # print(self.polygon)
                # b_xmin = int((box.getElementsByTagName(
                #     "xmin")[0]).childNodes[0].data)
                # b_ymin = int((box.getElementsByTagName(
                #     "ymin")[0]).childNodes[0].data)
                # b_xmax = int((box.getElementsByTagName(
                #     "xmax")[0]).childNodes[0].data)
                # b_ymax = int((box.getElementsByTagName(
                #     "ymax")[0]).childNodes[0].data)
                # self.name = dom.getElementsByTagName(
                #     'filename')[0].childNodes[0].data
                # print(self.name)
                # self.bbox.append([[b_xmin, b_ymin], [b_xmax, b_ymax]])
            # print(bbox)

            # print(self.polygon)
            return self.im_shape, self.bbox

    def xyxy2yolo(self):
        self.bboxyolo = []
        for bx in self.bbox:
            cent = (bx[1]+bx[0])/2
            dwh = bx[1]-bx[0]
            # self.bboxyolo.append([bx[0]/self.im_shape,dwh/self.im_shape])
            self.bboxyolo.append([cent/self.im_shape, dwh/self.im_shape])

        # for o_cls, pt in self.polygon:
        #     # print(o_cls[-5:])
        #     if o_cls[-5:]=="plate":
        #         # print(sorted(pt))
        #         plx=pt.min(0)
        #         prx=pt.max(0)
        #         self.bboxyolo.append([plx/self.im_shape,prx/self.im_shape])
        return np.array(self.bboxyolo).reshape(-1, 4), self.obj_name

    def imshow(self):
        img = cv2.imread(ROOT_PATH+"/"+self.name)
        for bx in self.bbox:
            cv2.rectangle(img, bx[0], bx[1], (0, 255, 0), 1, 1)
        # cv2.imshow("a", img)
        # cv2.waitKey()
        # img=cv2.cvtColor(img, cv2.COLOR_BGR2BGRA);
        for clas, point in self.polygon:
            # cv2.polylines(img, [point], 1, (0,0,0,255))
            # print(point)
            img = put_mask(img, point)
        size_decrease = (int(img.shape[1]/2), int(img.shape[0]/2))
        img_decrease = cv2.resize(
            img, size_decrease, interpolation=cv2.INTER_CUBIC)
        # mask = im
        # cv2.imshow("Mask", mask)
        # masked = cv2.bitwise_and(image, image, mask=mask)
        cv2.imshow("Mask to Image", img_decrease)
        cv2.waitKey(0)


# img = cv2.imread("d.jpg")
# img=cv2.resize(img,(960,540),cv2.INTER_CUBIC)
# cv2.imshow("Mask to Image", img)
# cv2.imwrite("a.jpg",img)
# cv2.waitKey(0)

categories2d = [
    {
        "supercategory": "none",
        "id": 0,
        "name": "open_bed_heavy_truck"
    },
    {
        "supercategory": "none",
        "id": 1,
        "name": "car"
    },
    {
        "supercategory": "none",
        "id": 2,
        "name": "bus"
    },
    {
        "supercategory": "none",
        "id": 3,
        "name": "van"
    },
    {
        "supercategory": "none",
        "id": 4,
        "name": "open_bed_light_truck"
    },
    {
        "supercategory": "none",
        "id": 5,
        "name": "others_truck"
    },
    {
        "supercategory": "none",
        "id": 6,
        "name": "close_bed_heavy_truck"
    },
    {
        "supercategory": "none",
        "id": 7,
        "name": "close_bed_light_truck"
    }
]
annotations = {
    "segmentation": [
        [
            510.66,
            423.01
        ]
    ],
    "area": 702.1057499999998,
    "iscrowd": 0,
    "image_id": 289343,
    "bbox": [
        473.07,
        395.93,
        38.65,
        28.67
    ],
    "category_id": 18,
    "id": 1768
}
images = {
    "license": 4,
    "file_name": "000000397133.jpg",
    "coco_url": "http://images.cocodataset.org/val2017/000000397133.jpg",
    "height": 427,
    "width": 640,
    "date_captured": "2013-11-14 17:02:52",
    "flickr_url": "http://farm7.staticflickr.com/6116/6255196340_da26cf2c9e_z.jpg",
    "id": 397133
}
licenses = {
    "url": "http://creativecommons.org/licenses/by-nc-sa/2.0/",
    "id": 1,
    "name": "Attribution-NonCommercial-ShareAlike License"
}
info = {
    "description": "COCO 2017 Dataset",
    "url": "http://cocodataset.org",
    "version": "1.0",
    "year": 2017,
    "contributor": "COCO Consortium",
    "date_created": "2017/09/01"
}
if __name__ == "__main__":

    xmls = glob.glob("%s/1441/*.xml" % ROOT_PATH)

    # ftxt = open("train.txt", mode="w")

    os.system("mkdir -p coco/annotations")
    os.system("mkdir -p coco/images")
    annotations_list, image_list = [], []
    img_cnt, anno_cnt = 0, 0

    for i, xml_path in enumerate(xmls):
        reader = XmlTester()
        im_shape, bbox = reader.load(xml_path)



            
        if len(reader.polygon) > 0:

            img_path = xml_path[:-3]+"jpg"
            img = cv2.imread(img_path)

            img_ = {
                "license": 0,
                "file_name": str(img_cnt)+".jpg",
                "coco_url": "",
                "height": img.shape[0],
                "width": img.shape[1],
                "date_captured": "2013-11-14 17:02:52",
                "flickr_url": "",
                "id": img_cnt,
            }


            # img_src_size=img.shape()
            # img = cv2.resize(img, (1920, 1080), interpolation=cv2.INTER_CUBIC)
            # cv2.imwrite("%s/coco/images/%d.jpg" % (ROOT_PATH, img_cnt),img)

            img_cnt+=1
            image_list.append(img_)
        else :
            print(i,xml_path)
        # print(reader.polygon)
        for cls, box in reader.polygon:
            print(cls,box)
            plx=box.min(0)
            prx=box.max(0)
            w,h=prx-plx
            # print(cls)
            # box=box/im_shape*np.array([1920,1080])
            # print(box,im_shape,k)
            # exit()
            # break
            annotations = {
                "segmentation": box.reshape(1,-1).tolist(),
                "area": 702.1057499999998,
                "iscrowd": 0,
                "image_id": img_cnt-1,
                "bbox": [
                    int(plx[0]),
                    int(plx[1]),
                    int(w),
                    int(h)
                ],
                "category_id": class_name.index(cls),
                "id": anno_cnt
            }
            anno_cnt+=1
            annotations_list.append(annotations)
            # print(type(img_))
            #


        # reader.imshow()
        # print(im_shape)

        # flb = open("%s/labels/%d.txt" % (ROOT_PATH, i), mode="w")

        #
        # ftxt.write("%s/images/%d.png\n" % (ROOT_PATH, i))
        # b, ns = reader.xyxy2yolo()
        # # print(b)
        # for b_, n in zip(b, ns):
        #     target = n if len(n) < 4 else n[-5:]
        #     # print(target,n)
        #     if class_name.count(target) != 0:
        #         indx = class_name.index(target)
        #         flb.write("%d %f %f %f %f\n" %
        #                   (indx, b_[0], b_[1], b_[2], b_[3]))

        # flb.close()
        # for o_cls, pt in reader.polygon:
        #     # print(o_cls[-5:])
        #     if o_cls[-5:]=="plate":
        #         print(o_cls, pt)
    # ftxt.close()
    userInfo = {"info": info, "licenses": licenses, "images": image_list,
                "annotations": annotations_list, "categories": categoriesmask}

    with open('coco/annotations/instances_train2017.json', 'w') as obj:
        json.dump(userInfo, obj)
    with open('coco/annotations/instances_val2017.json', 'w') as obj:
        json.dump(userInfo, obj)